import os
import pandas as pd

pd.set_option('display.max_columns', None)
pd.set_option('display.width', 200)


class Database(object):
    def __init__(self):
        self.path = ""
        self.raw_data = pd.DataFrame()
        self.path = os.path.dirname(__file__)


class eft_rt_data(Database):
    def __init__(self):
        super().__init__()
        self.path = os.path.join(self.path, "etf_min_rt_data", "2024_code")

    def load_data(self, code=["159150"], start_date=None, end_date=None):
        print("=========开始数据提取=========")
        # 校验时间
        if start_date > end_date:
            raise RuntimeError("start_date must be before end_date")

        for etf_code in code:
            file_list = os.listdir(self.path)
            etf_file = etf_code + ".csv"
            if etf_file in file_list:
                filename = os.path.join(self.path, etf_file)
                print("=========数据来源路径：%s=========" %filename)
                temp_df = pd.read_csv(filename)
                self.raw_data = pd.concat([self.raw_data, temp_df])
            else:
                raise RuntimeError(etf_code, " code is not exist")

        self.raw_data["日期"] = pd.to_datetime(self.raw_data["date"])
        self.raw_data["日期"] = self.raw_data["日期"].dt.strftime("%Y%m%d")

        if (start_date is not None) & (end_date is not None):
            print("date time is in {} - {}".format(start_date, end_date))
            self.raw_data = self.raw_data[(self.raw_data["日期"] >= start_date) & (self.raw_data["日期"] <= end_date)]

        print("数据行列数量 = ", self.raw_data.shape)
        print("数据字段为 = ", self.raw_data.columns)
        print("=========已完成数据提取=========")
        return self.raw_data

class etf_data(Database):
    def __init__(self):
        super().__init__()
        self.current_path = os.path.join(self.path, "etf_min_data")
        self.history_path = os.path.join(self.path, "etf_hist_data")

    def load_data(self, folder="5"):
        path = os.path.join(self.current_path, folder)
        for file in os.listdir(path):
            file_path = os.path.join(path, file)
            df = pd.read_csv(file_path)
            df["code"] = file.split(".")[0].split("_")[0]
            self.raw_data = pd.concat([self.raw_data, df])
        return self.raw_data


class etf_data_download(Database):
    def __init__(self):
        super().__init__()
        self.current_path = os.path.join(self.path, "2024")

    def load_data(self, code=None, start_date=None, end_date=None):
        print("=========开始数据提取=========")
        date_file_list = os.listdir(self.current_path)
        # 校验开始时间
        if start_date == None:
            start_date = min(date_file_list)
        else:
            if type(start_date) != str:
                raise TypeError("start_date type is error")
            elif start_date not in date_file_list:
                raise ValueError("start_date value error")

        # 校验结束时间
        if end_date == None:
            end_date = max(date_file_list)
        else:
            if type(end_date) != str:
                raise TypeError("end_date type is error")
            elif end_date not in date_file_list:
                raise ValueError("end_date value error")


        print("date time is in {} - {}".format(start_date, end_date))

        for date_filename in date_file_list:
            if date_filename >= start_date and date_filename <= end_date:
                date_folder = os.path.join(self.current_path, date_filename)
                for file_name in os.listdir(date_folder):
                    if code is not None:
                        code_name = file_name.split(".")[0].split("_")[0]
                        time = file_name.split(".")[0].split("_")[1]
                        if type(code) == str:
                            if code_name == code:
                                filename = os.path.join(date_folder, file_name)
                                temp_df = pd.read_csv(filename, encoding="gbk")
                                temp_df["filetime"] = time
                                temp_df["code"] = code_name
                                self.raw_data = pd.concat([self.raw_data, temp_df])
                        elif type(code) == list:
                            if code_name in code:
                                filename = os.path.join(date_folder, file_name)
                                temp_df = pd.read_csv(filename, encoding="gbk")
                                temp_df["filetime"] = time
                                temp_df["code"] = code_name
                                self.raw_data = pd.concat([self.raw_data, temp_df])
                        else:
                            raise TypeError("code type must be list or string")
                    else:

                        code_name = file_name.split(".")[0].split("_")[0]
                        time = file_name.split(".")[0].split("_")[1]
                        filename = os.path.join(date_folder, file_name)
                        temp_df = pd.read_csv(filename, encoding="gbk")
                        temp_df["filetime"] = time
                        temp_df["code"] = code_name
                        self.raw_data = pd.concat([self.raw_data, temp_df])
        if code is not None:
            print("已提取 {} 数据".format(code))
        else:
            print("已提取数据库全部数据")
        print("数据行列数量 = ", self.raw_data.shape)
        print("数据字段为 = ", self.raw_data.columns)
        print("=========已完成数据提取=========")
        print()
        return self.raw_data


if __name__ == "__main__":
    # etf_data = etf_data_download()
    # df = etf_data.load_data(code=["159001", "159003"], start_date="20240102", end_date="20240304")
    # print(df)

    data = eft_rt_data()
    df = data.load_data(code=["159001"], start_date="20240801", end_date="20240804")
